# -*- coding: utf-8 -*-
# created on 2016/5/10


import collections
from sympy import solveset
from itertools import product
from sympy.abc import z
from mathsolver.functions.base import *
from mathsolver.functions.root.stepEquations import Step2Equations, Step3Equations
from mathsolver.functions.sympy_utils import get_all_child


class JieFangChenZu(BaseFunction):
    def solver(self, *args):
        symbols = None
        if len(args) == 1:
            symbols = args[0].free_symbols()  # .intersection([x,y])
            if x in symbols and y in symbols and len(args[0].value) == 2:
                symbols = {x, y}

        elif len(args) == 2:
            symbols = args[1].sympify()

        eqs = args[0].sympify()
        free_symbols = []
        free_symbols_num = []
        for f1, f2 in eqs:
            free_symbol = (f1 - f2).free_symbols
            free_symbols.extend(list(free_symbol))
            free_symbols_num.append(len(free_symbol))

        simple_solve = []
        for f1, f2 in eqs:
            symbols_eq = (f1 - f2).free_symbols
            diff_symbols = symbols_eq.intersection(symbols)
            if len(diff_symbols) == 1:
                simple_solve.append(f1 - f2)

        if len(simple_solve) == len(eqs):
            answer = collections.defaultdict(list)
            pairs = []
            keys = []
            for eq in simple_solve:
                sym = list(eq.free_symbols.intersection(symbols))[0]
                pairs.append(list(solveset(eq, sym)))
                keys.append(sym)

            for values in itertools.product(*pairs):
                for k, value in zip(keys, values):
                    answer[k].append(value)

            self.output.append(BaseSymbolValue(answer))
            return self

        solver = None
        try:
            if len(eqs) == 2:
                solver = Step2Equations([Eq(*eqs[0]), Eq(*eqs[1])], symbols).solve()
            elif len(eqs) == 3:
                solver = Step3Equations([Eq(*eqs[0]), Eq(*eqs[1]), Eq(*eqs[2])]).solve()
        except Exception:
            pass

        if solver:
            answer = solver.answer
            steps = solver.printing()

            if not self.verbose and any([not f1.is_Symbol for f1, f2 in eqs]):
                self.steps.append(["", steps[-1][1]])
            elif self.verbose:
                self.steps = steps
                self.label.add("解方程组")
        else:
            sympy_answer = solve([Eq(*eq) for eq in eqs], symbols)
            if isinstance(sympy_answer, dict):
                answer = sympy_answer
            elif sympy_answer and isinstance(sympy_answer, list):
                answer = {}
                keys = sympy_answer[0].keys()
                for key in keys:
                    answer[key] = [ans[key] for ans in sympy_answer]

        known_sym, known_solve = zip(*answer.items())
        answer = collections.defaultdict(list)
        for known in zip(*known_solve):
            answer[known_sym].append(tuple(known))
        self.output.append(BaseSymbolValue(answer))
        return self


class FangChenZuQiaoSuan001(BaseFunction):
    def solver(self, *args):
        """
        已知方程组，求代数式的值。
        FangChenZuQiaoSuan001().solver(BaseEqs([['M', '3*(a)**(2)-2*a*b-4*(b)**(2)'], ['NN', '4*(a)**(2)+5*a*b-(b)**(2)']]), BasePoly("3*M-2*NN"))
        FangChenZuQiaoSuan001().solver(BaseEqs([['Abs(a)', '3'], ['Abs(b)', '2'], ['Abs(a/b)', 'a/b']]), BasePoly("3*a - 4*b"))
        :param args:
        :return:
        """
        eqs = args[0].sympify()
        poly = args[1].sympify()

        symbols = poly.free_symbols
        answer = dict([(sym, []) for sym in symbols])
        # |a|=3,|b|=2,|\frac{a}{b}|=\frac{a}{b},求3a-2b的值.
        simple_solve = [(i, eq) for i, eq in enumerate(eqs) if len(eq[0].free_symbols.intersection(symbols)) == 1]
        simple_index = [i for i, eq in simple_solve]
        assert len(simple_solve) > 0
        for i in simple_index:
            eq = Eq(*eqs[i])
            sym = eq.free_symbols.intersection(symbols).pop()
            if str(sym) == str(eq.args[0]):
                answer[sym].append(eq.args[1])
            else:
                answer[sym].extend(solveset(eq, sym, S.Reals))

        all_answer = list(product(*[v for k, v in answer.items()]))
        symbols = answer.keys()
        answers = [dict(zip(symbols, comb)) for comb in all_answer]
        rest_index = [i for i in range(len(eqs)) if i not in simple_index]
        for i in rest_index:
            f1, f2 = eqs[i]
            answers = [answer for answer in answers if f1.subs(answer) == f2.subs(answer)]

        assert len(answers) > 0
        for answer in answers:
            number = poly.subs(answer)
            self.steps.append(["将%s代入代数式，得：" % self.output_symbol_value(answer),
                               "%s = %s" % (new_latex(poly), new_latex(number))])
            self.output.append(BaseNumber(number))

        return self


class FangChenZuQiaoSuan002(BaseFunction):
    @staticmethod
    def is_solvable(temp_poly, candidates):
        targets = []
        for i in range(3):
            for cand in candidates:
                if temp_poly.has(cand):
                    temp_poly = temp_poly.subs({cand: 1})
                    targets.append(candidates.index(cand))
                    break
                elif temp_poly.is_Add:
                    a, h = temp_poly.as_independent(*temp_poly.free_symbols)
                    if len((h / cand).cancel().free_symbols) == 0:
                        temp_poly = (h / cand).cancel()
                        targets.append(candidates.index(cand))
            if len(temp_poly.free_symbols) == 0:
                return targets
        return []

    def solver(self, *args):
        """
        FangChenZuQiaoSuan002().solver(BaseEqs([['(3)**(x)', 'a'], ['(3)**(y)', 'b']]), BasePoly("(3)**(x-y)"))
        FangChenZuQiaoSuan002().solver(BaseEqs([['a**x', '2'], ['a**y', '3']]), BasePoly("a**(2*x+3*y)"))
        FangChenZuQiaoSuan002().solver(BaseEqs([['a -b', '1'], ['a*b', '-2']]), BasePoly("(a+1)*(b-1)"))
        FangChenZuQiaoSuan002().solver(BaseEqs([['(x+y)**2', '5'], ['(x-y)**2', '3']]), BasePoly("x*y"))
        FangChenZuQiaoSuan002().solver(BaseEqs([['3**m', '2'], ['3**n', '4']]), BasePoly("3**(m+n+1)"))
        :param args:
        :return:
        """
        eqs = args[0].sympify()
        poly = args[1].sympify()

        symbols = poly.free_symbols
        new_eqs = []
        for f1, f2 in eqs:
            eq = f1 - f2
            if eq.is_Add:
                a, h = eq.as_independent(*symbols)
                new_eqs.append([h, -a])
            else:
                new_eqs.append([f1, f2])

        temp = [str(eq[0]) for eq in new_eqs]
        if str(poly) in temp:
            index = temp.index(str(poly))
            value = new_eqs[index][1]
            self.steps.append(["", "%s的值为%s" % (new_latex(poly), new_latex(value))])
            self.output.append(BaseNumber(value))

        # 若3^{m}=2,3^{n}=4,则3^{m+n-1}=..
        elif poly.is_Pow and (len(poly.args[0].free_symbols) == 0 or poly.args[0].is_Symbol):
            _base = poly.args[0]
            if all([f1.is_Pow and f1.args[0] == _base for f1, _ in new_eqs]):
                dairu_value = dict([(f1.args[1], f2) for f1, f2 in new_eqs])
                value = poly.args[1].subs(dairu_value)

                # ratio = (poly/new_eqs[0][0]/new_eqs[1][0]).cancel()
                # print ratio
                symbol_value = dict([(f1, f2) for f1, f2 in new_eqs])
                out = Pow(_base, value)
                self.steps.append(["将%s代入代数式，得：" % self.output_symbol_value(symbol_value),
                                   "%s = %s" % (new_latex(poly), new_latex(out))])
                self.output.append(BasePoly(out))
                return self

        elif len(symbols) == 2 and not get_all_child(poly, lambda xx: xx.is_Pow and abs(xx.args[1]) > 5):
            s1, s2 = list(symbols)
            candidates = [s1 * s2,  # 0
                          (s1 + s2) ** 2,  # 1
                          (s1 - s2) ** 2,  # 2
                          (s2 - s1) ** 2,  # 3
                          s1 ** 2 + s2 ** 2,  # 4
                          s1 ** 2 - s2 ** 2,  # 5
                          s2 ** 2 - s1 ** 2,  # 6
                          s1 + s2,  # 7
                          s1 - s2,  # 8
                          s2 - s1,  # 9
                          1 / s1 + 1 / s2  # 10
                          ]

            eq_dict = {(0, 1, 2): Eq(4 * x, y - z), (0, 1, 3): Eq(4 * x, y - z), (0, 1, 4): Eq(2 * x, y - z),
                       (0, 1, 8): Eq(4 * x, y - z ** 2), (0, 1, 9): Eq(4 * x, y - z ** 2),
                       (0, 2, 4): Eq(2 * x, z - y), (0, 2, 7): Eq(4 * x, z ** 2 - y),
                       (0, 3, 4): Eq(2 * x, z - y), (0, 3, 7): Eq(4 * x, z ** 2 - y),
                       (0, 4, 8): Eq(2 * x, y - z ** 2), (0, 4, 9): Eq(2 * x, y - z ** 2),
                       (0, 4, 7): Eq(2 * x, z ** 2 - y),
                       (1, 7): Eq(x, y ** 2), (2, 8): Eq(x, y ** 2), (2, 9): Eq(x, y ** 2), (8, 9): Eq(x, -y),
                       (5, 4): Eq(x, -y),
                       (3, 8): Eq(x, y ** 2), (3, 9): Eq(x, y ** 2), (5, 7, 8): Eq(x, y * z), (5, 7, 9): Eq(x, y * z),
                       (6, 7, 8): Eq(x, -y * z), (6, 7, 9): Eq(x, -y * z),
                       (0, 7, 10): Eq(y / x, z)}

            if all([f1 in candidates for f1, _ in new_eqs]):
                known = dict([(candidates.index(f1), f2) for f1, f2 in new_eqs])
                target = self.is_solvable(poly, candidates)
                expr = poly
                if not target:
                    target = self.is_solvable(poly.expand(), candidates)
                    expr = poly.expand()

                if not target:
                    target = self.is_solvable(poly.expand().factor(), candidates)
                    expr = poly.expand().factor()

                target_value = {}

                for tar in target:
                    if tar in known:
                        target_value[candidates[tar]] = known[tar]
                        continue

                    all_index = sorted([tar] + known.keys())
                    select_index = [k for k in eq_dict if set(k).intersection(all_index) == set(k) and tar in k][0]
                    symbol_names = dict(
                        [([x, y, z][i1], "(%s)" % new_latex(candidates[i2])) for i1, i2 in enumerate(select_index)])
                    variables = dict([(index, v) for index, v in zip(select_index, [x, y, z])])
                    ans = solveset(eq_dict[select_index], variables[tar]).args[0]
                    value = ans.subs(
                        dict([(variables[key], value) for key, value in known.items() if key in variables]))
                    target_value[candidates[tar]] = value
                    self.steps.append(["", "%s = %s = %s " % (
                        new_latex(candidates[tar]), new_latex(ans, symbol_names=symbol_names), new_latex(value))])

                if len(target_value) > 1 or str(list(target_value.keys())[0]) != str(expr):
                    # dairu
                    for i in range(3):
                        expr = expr.subs(target_value)
                        if len(expr.free_symbols) > 0 and expr.is_Add:
                            a, h = expr.as_independent(*expr.free_symbols)
                            for k, v in target_value.items():
                                if len((h / k).cancel().free_symbols) == 0:
                                    expr = (h / k).cancel() * v + a
                                    continue
                        else:
                            continue
                    assert len(expr.free_symbols) == 0
                    self.steps.append(["将%s代入代数式，得：" % self.output_symbol_value(target_value),
                                       "%s = %s" % (new_latex(poly), new_latex(expr))])
                    self.output.append(BaseNumber(expr))
                else:
                    self.output.append(BaseNumber(list(target_value.values())[0]))
        assert len(self.output) > 0
        return self


if __name__ == '__main__':
    pass
